// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


// Adapted from Apache Arrow.

#include <pollux/dwio/parquet/writer/arrow/tests/ColumnScanner.h>

#include <cstdint>
#include <memory>

#include <pollux/dwio/parquet/writer/arrow/tests/ColumnReader.h>

using arrow::MemoryPool;

namespace kumo::pollux::parquet::arrow {

std::shared_ptr<Scanner> Scanner::Make(
    std::shared_ptr<ColumnReader> col_reader,
    int64_t batch_size,
    MemoryPool* pool) {
  switch (col_reader->type()) {
    case Type::BOOLEAN:
      return std::make_shared<BoolScanner>(
          std::move(col_reader), batch_size, pool);
    case Type::INT32:
      return std::make_shared<Int32Scanner>(
          std::move(col_reader), batch_size, pool);
    case Type::INT64:
      return std::make_shared<Int64Scanner>(
          std::move(col_reader), batch_size, pool);
    case Type::INT96:
      return std::make_shared<Int96Scanner>(
          std::move(col_reader), batch_size, pool);
    case Type::FLOAT:
      return std::make_shared<FloatScanner>(
          std::move(col_reader), batch_size, pool);
    case Type::DOUBLE:
      return std::make_shared<DoubleScanner>(
          std::move(col_reader), batch_size, pool);
    case Type::BYTE_ARRAY:
      return std::make_shared<ByteArrayScanner>(
          std::move(col_reader), batch_size, pool);
    case Type::FIXED_LEN_BYTE_ARRAY:
      return std::make_shared<FixedLenByteArrayScanner>(
          std::move(col_reader), batch_size, pool);
    default:
      ParquetException::NYI("type reader not implemented");
  }
  // Unreachable code, but suppress compiler warning
  return std::shared_ptr<Scanner>(nullptr);
}

int64_t ScanAllValues(
    int32_t batch_size,
    int16_t* def_levels,
    int16_t* rep_levels,
    uint8_t* values,
    int64_t* values_buffered,
    ColumnReader* reader) {
  switch (reader->type()) {
    case parquet::Type::BOOLEAN:
      return ScanAll<BoolReader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::INT32:
      return ScanAll<Int32Reader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::INT64:
      return ScanAll<Int64Reader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::INT96:
      return ScanAll<Int96Reader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::FLOAT:
      return ScanAll<FloatReader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::DOUBLE:
      return ScanAll<DoubleReader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::BYTE_ARRAY:
      return ScanAll<ByteArrayReader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    case parquet::Type::FIXED_LEN_BYTE_ARRAY:
      return ScanAll<FixedLenByteArrayReader>(
          batch_size, def_levels, rep_levels, values, values_buffered, reader);
    default:
      ParquetException::NYI("type reader not implemented");
  }
  // Unreachable code, but suppress compiler warning
  return 0;
}

} // namespace kumo::pollux::parquet::arrow
